import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

# 生成随机城市坐标
num_cities = 10
cities = np.random.rand(num_cities, 2)


# 计算城市之间的距离矩阵
def calc_distance_matrix(cities):
    num_cities = len(cities)
    distance_matrix = np.zeros((num_cities, num_cities))
    for i in range(num_cities):
        for j in range(num_cities):
            distance_matrix[i, j] = np.linalg.norm(cities[i] - cities[j])  # 计算两城市之间的欧几里得距离
    return distance_matrix


# 生成距离矩阵
distance_matrix = calc_distance_matrix(cities)


# 最近邻算法求解TSP（旅行商问题）
def nearest_neighbor_algorithm(distance_matrix):
    num_cities = len(distance_matrix)
    visited = [False] * num_cities  # 记录每个城市是否已被访问
    path = [0]  # 初始化路径，从第一个城市开始
    visited[0] = True  # 标记第一个城市为已访问
    for _ in range(num_cities - 1):
        last_city = path[-1]
        # 查找距离当前城市最近且未访问的城市
        next_city = np.argmin([distance_matrix[last_city, j] if not visited[j] else np.inf for j in range(num_cities)])
        path.append(next_city)  # 将下一个城市添加到路径中
        visited[next_city] = True  # 标记下一个城市为已访问
    path.append(0)  # 回到起始城市，完成环路
    return path


# 生成旅行路径
path = nearest_neighbor_algorithm(distance_matrix)


# 验证路径完整性
def verify_path(path, num_cities):
    if len(path) != num_cities + 1:  # 路径长度应为城市数量加一（返回起始点）
        return False
    if len(set(path[:-1])) != num_cities:  # 检查路径中是否包含所有城市
        return False
    if path[0] != path[-1]:  # 检查路径是否以起始城市开始并结束
        return False
    return True


# 计算路径总距离
def calculate_total_distance(path, distance_matrix):
    total_distance = 0
    for i in range(len(path) - 1):
        total_distance += distance_matrix[path[i], path[i + 1]]  # 累加每段路径的距离
    return total_distance


# 验证路径
is_valid_path = verify_path(path, num_cities)
total_distance = calculate_total_distance(path, distance_matrix)
validation_text = "Path validation: PASSED" if is_valid_path else "Path validation: FAILED"
distance_text = f"Total distance: {total_distance:.2f}"

# 创建动画
fig, ax = plt.subplots()
scat = ax.scatter(cities[:, 0], cities[:, 1], c='blue')  # 绘制城市坐标点
line, = ax.plot([], [], 'r-', lw=2)  # 初始化路径线
step_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)  # 初始化步骤文本

# 在图形外部显示验证信息和路径总距离
fig.text(0.02, 0.95, validation_text, transform=fig.transFigure)
fig.text(0.02, 0.92, distance_text, transform=fig.transFigure)

# 标出起始城市
start_city = cities[0]
ax.scatter(start_city[0], start_city[1], c='green', s=200)  # 起始城市用绿色标记


# # 显示每个城市的坐标
# for i, (x, y) in enumerate(cities):
#     ax.text(x, y, f'({x:.2f}, {y:.2f})', fontsize=9, ha='right')


# 初始化动画帧
def init():
    line.set_data([], [])  # 清空路径线
    step_text.set_text('')  # 清空步骤文本
    return line, step_text


# 更新动画帧
def update(frame):
    step = frame + 1
    current_path = path[:step]  # 获取当前路径
    x = [cities[i][0] for i in current_path]
    y = [cities[i][1] for i in current_path]
    line.set_data(x, y)  # 更新路径线
    step_text.set_text(f'Step: {step}')  # 更新步骤文本
    return line, step_text


# 创建动画
ani = animation.FuncAnimation(fig, update, frames=len(path) - 1, init_func=init, blit=True, repeat=False, interval=500)
plt.show()
